# -*- coding: utf-8 -*-
"""
Created on Thu Apr 18 15:05:23 2019

@author: david
"""
import datetime
import time

import matplotlib.pyplot as plt
import numpy as np
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False    # 解决负号显示问题


'''
改变问题规模时需要改动的参数：工件参数workpiece，工序数目process，机器数目machine，
机器选择部分的范围上限，迭代次数（规模小的时候可能降低迭代次数）
'''

# 读取原始数据
contents = []
with open(r"C:\Users\L\Downloads\FJSP-master\data\data_first.txt") as f:
    string = f.readlines()
    for item in string:
        contents.append(item.strip().split(" "))


# 对工序部分进行处理
def handle(x):
    """
    处理工序列表，统计每个工序的出现次数并记录。

    输入：粒子的位置，输出：对工序部分处理后的列表
    """
    # 初始化一个长度为workpiece的数组，用于统计每个工序的出现次数
    piece_mark = np.zeros(workpiece)  # 统计工序的标志

    # 初始化一个空列表，用于存放经过处理后的工序信息
    array = []  # 经过处理后的工序列表

    # 遍历粒子的位置，处理每个工序
    for i in range(total_process):
        # 根据粒子位置更新工序计数
        piece_mark[int(x[i] - 1)] += 1

        # 将当前工序及其出现次数作为一个元组添加到列表中
        array.append((int(x[i]), int(piece_mark[int(x[i] - 1)])))

    # 返回处理后的工序列表
    return array



def initpopvfit():
    """
    初始化种群、速度和适应度值。

    初始化一个种群矩阵pop，其中每行代表一个粒子，粒子的前半部分表示工件的加工顺序，
    后半部分表示每道工序分配的机器。初始化一个速度矩阵v，用于更新粒子的位置。
    初始化一个适应度数组fitness，用于存储每个粒子的适应度值。

    返回:
        pop: 种群矩阵
        v: 速度矩阵
        fitness: 适应度数组
    """
    # 初始化种群矩阵和速度矩阵以及适应度数组
    pop = np.zeros((popsize, total_process * 2))
    v = np.zeros((popsize, total_process * 2))
    fitness = np.zeros(popsize)

    # 为种群中的每个粒子进行初始化
    for i in range(popsize):
        # 初始化工序部分
        for j in range(workpiece):
            for p in range(process):
                pop[i][j * process + p] = j + 1
        np.random.shuffle(pop[i][:total_process])

        # 初始化机器部分
        for j in range(total_process):
            index = np.random.randint(0, machine)
            while contents[j][index] == "-":
                index = np.random.randint(0, machine)
            pop[i][j + total_process] = index + 1

        # 计算各粒子初始的适应度
        fitness[i] = calculate(pop[i])
    return pop, v, fitness



def calculate(x):
    """
    计算粒子的适应度值。

    输入: 粒子位置
    输出: 粒子适应度值
    """
    # 初始化每个机器上的完工时间为0
    Tm = np.zeros(machine)
    # 初始化每个工序的完成时间为0
    Te = np.zeros((workpiece, process))
    # 对粒子位置进行处理，以适应后续的计算
    array = handle(x)

    # 遍历每一个工序
    for i in range(total_process):
        # 计算确定机器的索引
        machine_index = int(x[total_process + (array[i][0] - 1) * process + (array[i][1] - 1)]) - 1
        # 计算确定工序的索引
        process_index = (array[i][0] - 1) * process + (array[i][1] - 1)
        # 获取当前工序的加工时间
        process_time = int(contents[process_index][machine_index])

        # 如果当前工序是第一个工序，则直接更新机器的完工时间和工序的完成时间
        if array[i][1] == 1:
            Tm[machine_index] += process_time
            Te[array[i][0] - 1][array[i][1] - 1] = Tm[machine_index]
        else:
            # 如果不是第一个工序，则需要考虑前一道工序的完成时间
            Tm[machine_index] = max(Te[array[i][0] - 1][array[i][1] - 2], Tm[machine_index]) + process_time
            Te[array[i][0] - 1][array[i][1] - 1] = Tm[machine_index]

    # 返回所有机器中最长的完工时间作为粒子的适应度值
    return max(Tm)



def getinitbest(fitness, pop):
    """
    初始化粒子群中的全局最优和个体最优位置及其适应度值

    该函数通过比较粒子群中所有粒子的适应度值，找出适应度值最小的粒子作为全局最优，
    同时将每个粒子的位置作为其个体最优位置初始化

    参数:
    fitness: 一维数组，包含粒子群中每个粒子的适应度值
    pop: 二维数组，包含粒子群中每个粒子的位置信息

    返回:
    gbestpop: 全局最优粒子的位置
    gbestfitness: 全局最优粒子的适应度值
    pbestpop: 二维数组，包含每个粒子的个体最优位置
    pbestfitness: 一维数组，包含每个粒子的个体最优适应度值
    """
    # 群体最优的粒子位置及其适应度值
    gbestpop, gbestfitness = pop[fitness.argmin()].copy(), fitness.min()
    # 个体最优的粒子位置及其适应度值,使用copy()使得对pop的改变不影响pbestpop，pbestfitness类似
    pbestpop, pbestfitness = pop.copy(), fitness.copy()
    return gbestpop, gbestfitness, pbestpop, pbestfitness



if __name__ == "__main__":
    # 初始化参数
    workpiece = 10  # 工件数目
    process = 5  # 每个工件的工序数目
    total_process = workpiece * process  # 工序的总数
    machine = 6  # 机器数目
    maxgen = 500  # 最大迭代次数
    w = 0.9  # 惯性权重
    lr = (2, 2)  # 加速因子
    popsize = 50  # 种群规模
    rangepop = (1, 6)  # 粒子编码中机器选择部分的范围
    clean_contents = []

    # 预处理数据，将工件在不同机器上的加工时间转换为概率分布
    for i in range(total_process):
        clean_contents.append([[int(contents[i][j]), j + 1] for j in range(machine) if contents[i][j] != "-"])
        temp_sum = 0
        for j in range(len(clean_contents[i])):
            temp_sum += 1 / clean_contents[i][j][0]
        for j in range(len(clean_contents[i])):
            clean_contents[i][j][0] = (1 / clean_contents[i][j][0]) / temp_sum
        clean_contents[i].sort()
        cumulation = 0
        for j in range(len(clean_contents[i])):
            cumulation += clean_contents[i][j][0]
            clean_contents[i][j][0] = cumulation

    # 初始化种群、速度和适应度
    pop, v, fitness = initpopvfit()
    # 获取初始最优解
    gbestpop, gbestfitness, pbestpop, pbestfitness = getinitbest(fitness, pop)

    # 记录每一代的迭代过程和全局最优解
    iter_process = np.zeros(maxgen)
    pso_base = np.zeros(maxgen)

    # 开始计时
    begin = time.time()
    # 主循环，进行迭代
    for i in range(maxgen):
        # 更新速度和位置
        for j in range(popsize):
            v[j] = w * v[j] + lr[0] * np.random.rand() * (pbestpop[j] - pop[j]) + lr[1] * np.random.rand() * (
                        gbestpop - pop[j])

        # 更新粒子位置，包括工序和机器两部分
        for j in range(popsize):
            store = []
            before = pop[j][:total_process].copy()
            pop[j] += v[j]
            reference = v[j][:total_process].copy()
            for p in range(total_process):
                store.append((reference[p], before[p]))
            store.sort()
            for p in range(total_process):
                pop[j][p] = store[p][1]

        pop = np.ceil(pop)

        for j in range(popsize):
            array = handle(pop[j])
            for p in range(total_process):
                if (pop[j][total_process + (array[p][0] - 1) * process + (array[p][1] - 1)] < rangepop[0] or pop[j][
                    total_process + (array[p][0] - 1) * process + (array[p][1] - 1)] > rangepop[1]) \
                        or (contents[(array[p][0] - 1) * process + (array[p][1] - 1)][int(
                    pop[j][total_process + (array[p][0] - 1) * process + (array[p][1] - 1)] - 1)] == "-"):
                    row = (array[p][0] - 1) * process + (array[p][1] - 1)
                    pop[j][total_process + (array[p][0] - 1) * process + (array[p][1] - 1)] = \
                    clean_contents[row][len(clean_contents[row]) - 1][1]

        # 记录当前迭代的最小适应度和全局最优解
        iter_process[i] = fitness.min()
        pso_base[i] = gbestfitness

        # 更新适应度
        for j in range(popsize):
            fitness[j] = calculate(pop[j])

        # 更新个体最优解
        for j in range(popsize):
            if fitness[j] < pbestfitness[j]:
                pbestfitness[j] = fitness[j]
                pbestpop[j] = pop[j].copy()

        # 更新全局最优解
        if pbestfitness.min() < gbestfitness:
            gbestfitness = pbestfitness.min()
            gbestpop = pop[pbestfitness.argmin()].copy()

    # 输出结果
    print("按照完全随机初始化的pso算法求得的最好的最大完工时间：", min(pso_base))
    print("按照完全随机初始化的pso算法求得的最好的工艺方案：", gbestpop)
    end = time.time()
    print("整个迭代过程所耗用的时间：{:.2f}s".format(end - begin))

    # 绘制迭代过程和全局最优解的变化情况
    fig = plt.figure(figsize=(10, 5))
    ax1 = fig.add_subplot(121)
    ax1.set_title("全局最优解的变化情况")
    ax1.plot(pso_base)
    ax2 = fig.add_subplot(122)
    ax2.set_title("每次迭代后种群适应度最小值的变化情况")
    ax2.plot(iter_process)
    # 按照当前时间保存图片，年月日用_隔开
    plt.savefig(str(datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')) + ".png")
    # plt.show()

